import logging
import os
import sys

from absl import app, flags

import numpy
import pandas

from coin.exchange.base.order_gateway import OrderSide
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.tool.signal_plot_base import dump_book

import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.tool.orderlogic_base as olbase
from coin.strategy.mm.tool.lm_signal_base import (dump_lm_signal, get_lm_product)

FLAGS = flags.FLAGS


def bisect_batch(lhs, rhs):
  i = 0
  j = 0
  rindex = []
  while i < len(lhs):
    while j < len(rhs):
      if lhs[i] > rhs[j]:
        j += 1
      else:
        break
    rindex.append(j)
    i += 1
  return rindex


class SignalAlignJob(object):
  def __init__(self, trading_date, time_range, profile):
    self.trading_date = trading_date
    self.time_range = time_range
    self.profile = profile
    self.product = get_lm_product(self.profile, self.trading_date)
    if not os.path.exists(self.root_dir):
      os.makedirs(self.root_dir)

  @property
  def root_dir(self):
    return "signal_plot"

  @property
  def filename(self):
    return os.path.join(
        self.root_dir,
        "%s.h5" % "_".join((self.trading_date.strftime("%Y%m%d"), self.profile, self.time_range)))

  def serialize(self):
    bbos = dump_book(self.product, self.trading_date, self.time_range)
    signals = dump_lm_signal(self.profile, self.trading_date, self.time_range, 0)
    # assert len(signals) > 0
    pandas.DataFrame(signals).to_hdf(self.filename, key='signals', mode='w')
    pandas.DataFrame(bbos).to_hdf(self.filename, key='bbos', mode='a')

  def deserialize(self):
    if not os.path.exists(self.filename):
      print("%s doesn't exist, dumping.." % self.filename)
      self.serialize()
    signals = pandas.read_hdf(self.filename, 'signals', mode='r')
    bbos = pandas.read_hdf(self.filename, 'bbos', mode='r')
    return signals, bbos


def plot_align_signal(signals, bbos, is_pass, threshold_bps, lean_bps, stacks):
  import matplotlib
  matplotlib.use("Agg")
  import matplotlib.pyplot as plt

  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.family'] = ['Dejavu Sans Mono']
  plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
  plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  signal_df = pandas.DataFrame(signals)

  signal_df['sign'] = ((signal_df['order_side'] == OrderSide.BUY) - 0.5) * 2

  stks = stacks

  if lean_bps.find("-") >= 0:
    open_lean_bps, close_lean_bps = lean_bps.split("-")
    open_lean_bps = float(open_lean_bps)
    close_lean_bps = float(close_lean_bps)
  else:
    open_lean_bps = float(lean_bps)
    close_lean_bps = float(lean_bps)

  chosen = signal_df['signal'] >= signal_df['order_price'] * (threshold_bps
                                                              - close_lean_bps * 0.5) * 1e-4

  signal_df = signal_df.loc[chosen].reset_index(drop=True)

  signs = signal_df['sign']
  stimes = signal_df['timestamp']
  signals_bps = signal_df['signal'] / signal_df['order_price'] * 1e4

  for stk in stks:
    pos = 0
    poss = []
    prev_action_time = 0
    for sign, signal_bps, stime in zip(signs, signals_bps, stimes):
      pos_ratio = pos / (stk / 2)
      open_abs_lean = abs(pos_ratio * open_lean_bps)
      close_abs_lean = abs(pos_ratio * close_lean_bps)
      if (sign > 0) == (pos_ratio > 0):
        adj_threshold_bps = threshold_bps + open_abs_lean
      else:
        adj_threshold_bps = threshold_bps - close_abs_lean

      if (stime - prev_action_time > flags.FLAGS.min_order_interval_ms * 1e6
          and signal_bps > adj_threshold_bps):
        if sign > 0:
          pos += 1
          prev_action_time = stime
        else:
          pos -= 1
          prev_action_time = stime
      pos = min(max(pos, -int(stk / 2)), int(stk / 2))
      poss.append(pos)

    signal_df['pos_%s' % stk] = poss

  signal_df = olbase.add_rets(signal_df, 'order_price')

  inv_aggps = signal_df['agg_price'][::-1]
  inv_ordps = signal_df['order_price'][::-1]
  inv_tss = signal_df['timestamp'][::-1]

  optis = [0.0, 1.0]

  for stk in stks:
    for optimism in optis:
      next_mtmp = None
      next_ts = None
      pnls = []
      tovs = []
      inv_poss = signal_df['pos_%s' % stk][::-1]
      prev_inv_pos = 0
      for inv_aggp, inv_ordp, inv_pos, inv_ts in zip(inv_aggps, inv_ordps, inv_poss, inv_tss):
        new_mtmp = optimism * inv_aggp + (1 - optimism) * inv_ordp
        ret = 0 if next_mtmp is None else next_mtmp - new_mtmp
        next_mtmp = new_mtmp
        if next_ts is not None and next_ts - inv_ts > 600 * 1e9:
          ret = 0

        pnls.append(ret * inv_pos)
        tovs.append(numpy.abs(prev_inv_pos - inv_pos) * inv_ordp)
        prev_inv_pos = inv_pos
        next_ts = inv_ts
      signal_df['pnl_opti%.1f_stk%s' % (optimism, stk)] = pnls[::-1]
      signal_df['tov_opti%.1f_stk%s' % (optimism, stk)] = tovs[::-1]

  signal_ts = pandas.DatetimeIndex(signal_df['timestamp'])
  signal_vals = signal_df['signal']
  order_sides = signal_df['order_side']
  order_prices = signal_df['order_price']

  bbodf = pandas.DataFrame(bbos)
  bookdf = bbodf.loc[~numpy.isnan(bbodf['ask0_price'])]
  trddf = bbodf.loc[~numpy.isnan(bbodf['trade_price'])]

  bbo_ts = pandas.DatetimeIndex(bookdf['timestamp'])
  ask0p = bookdf['ask0_price']
  bid0p = bookdf['bid0_price']

  trd_ts = pandas.DatetimeIndex(trddf['timestamp'])
  trdp = trddf['trade_price']

  plt.subplot(311)
  plt.plot(
      bbo_ts,
      ask0p,
      'k-',
      bbo_ts,
      bid0p,
      'k-',
      # trd_ts, trdp, 'kx',
      signal_ts[order_sides == OrderSide.BUY],
      order_prices[order_sides == OrderSide.BUY],
      'g.',
      signal_ts[order_sides == OrderSide.SELL],
      order_prices[order_sides == OrderSide.SELL],
      'r.',
      lw=0.5,
      markersize=1)
  plt.title("%s threshold_bps:%s, min_order_interval_ms: %.2f, signal plot" %
            (flags.FLAGS.profile, threshold_bps, flags.FLAGS.min_order_interval_ms))
  plt.subplot(312)

  # import pdb; pdb.set_trace()
  def get_ret_bps(signal_df, postfix):
    return (signal_df['pnl_%s' % postfix].sum() / signal_df['tov_%s' % postfix].sum()) * 1e4

  # def get_daysharpe()

  lstys = []
  postfixs = []
  pnlcols = []
  pnllegends = []
  posplots = []
  poslegends = []

  colarr = ['r', 'g', 'b', 'k', 'y']

  pnlrows = []

  for opti, lty in [(optis[0], '-'), (optis[1], '--')]:
    for i, stk in enumerate(stks):
      col = colarr[i]
      postfix = 'opti%.1f_stk%s' % (opti, stk)
      postfixs.append(postfix)
      pnlcols.append('pnl_%s' % postfix)
      lstys.append(col + lty)

      null_idx = numpy.diff(signal_ts).astype(int) > 1e11
      plot_graph = signal_df['pnl_%s' % postfix].cumsum()
      plot_graph[numpy.where(null_idx)[0]] = numpy.nan

      rets = numpy.nan_to_num(numpy.diff(plot_graph))
      tottime = numpy.diff(signal_ts)[~null_idx].sum()
      totsharpe = numpy.mean(rets) / numpy.std(rets)
      sqrtndays = numpy.sqrt(tottime / numpy.timedelta64(1, 'D'))
      daysharpe = totsharpe / sqrtndays

      plt.plot(signal_ts, plot_graph, col + lty, lw=0.5)
      pnllegends.append('%15s: %7.2f retbps, %1.3f daysharpe' %
                        (postfix, get_ret_bps(signal_df, postfix), daysharpe))
      pnlrows.append((opti,
                      stk,
                      get_ret_bps(signal_df, postfix),
                      daysharpe,
                      signal_df['pnl_%s' % postfix].sum()))
      if lty == '--':
        posplots.append((signal_ts, signal_df['pos_%s' % stk], col + ':'))
        poslegends.append('stk%s' % stk)
  plt.title("pnl plot per settings")
  plt.legend(pnllegends)

  plt.subplot(313)
  for posplot in posplots:
    plt.plot(*posplot, lw=0.2)
  plt.legend(poslegends)
  plt.title("pos plot per settings")

  msstr = str(int(FLAGS.min_order_interval_ms)).zfill(4)
  plt.savefig(os.path.join(
      "signal_plot",
      "%s_%s_th%s_ln%s_mo%s.png" %
      (FLAGS.trading_date, FLAGS.profile, threshold_bps, lean_bps, msstr)),
              dpi=400)
  plt.close()

  plt.rcParams['figure.figsize'] = 4, 3
  plt.rcParams['font.size'] = 5
  plt.rcParams['legend.fontsize'] = 4
  plt.rcParams['xtick.labelsize'] = 5
  plt.rcParams['ytick.labelsize'] = 5

  n = len(pandas.DataFrame(pnlrows).groupby(0))
  for i, pnldf in pandas.DataFrame(pnlrows).groupby(0):
    plt.subplot(n, 2, (i * 2) + 1)
    plt.plot(pnldf[1], pnldf[2], 'g.-', lw=0.5, markersize=1)
    plt.title("%s th%s ln%s" % (FLAGS.profile, threshold_bps, lean_bps))
    plt.xlabel('#stack')
    plt.legend(['retbps'], loc='upper left')
    plt.twinx()
    plt.plot(pnldf[1], pnldf[4], 'g.-', lw=0.5, markersize=1)
    plt.legend(['pnl'], loc='upper right')
    plt.subplot(n, 2, (i * 2) + 2)
    plt.plot(pnldf[1], pnldf[3], 'r.-', lw=0.5, markersize=1)
    plt.xlabel('#stack')
    plt.legend(['daysharpe'], loc='upper left')
  plt.tight_layout()
  plt.savefig(os.path.join(
      "signal_plot",
      "summ_%s_%s_th%s_ln%s_mo%s.png" %
      (FLAGS.trading_date, FLAGS.profile, threshold_bps, lean_bps, msstr)),
              dpi=200)
  plt.close()


def dump_signal_book(trading_date, time_range, profile):
  align_job = SignalAlignJob(trading_date, time_range, profile)
  return align_job.deserialize()


def launch_multi(trading_dates, time_range, profile, thresholds_bps, leans_bps, stacks):
  if len(trading_dates) == 1:
    signals, bbos = dump_signal_book(trading_dates[0], time_range, profile)
  else:
    from concurrent.futures import ProcessPoolExecutor
    futures = []
    with ProcessPoolExecutor(max_workers=5) as executor:
      for trading_date in trading_dates:
        futures.append(executor.submit(dump_signal_book, trading_date, time_range, profile))
    signals = []
    bbos = []
    for future in futures:
      asignals, abbos = future.result()
      signals.append(asignals)
      bbos.append(abbos)
    signals = pandas.concat(signals)
    bbos = pandas.concat(bbos)
  for threshold_bps, lean_bps in zip(thresholds_bps, leans_bps):
    plot_align_signal(signals,
                      bbos,
                      is_pass=('pass' in profile),
                      threshold_bps=float(threshold_bps),
                      lean_bps=lean_bps,
                      stacks=stacks)


def main(_):
  # import warnings; warnings.filterwarnings("error")
  trading_date = FLAGS.trading_date
  profile = FLAGS.profile
  time_range = FLAGS.time_range
  assert trading_date, '--trading_date must be specified.'
  assert profile, '--profile must be specified.'
  print('Running for %s %s %s moi %s...' %
        (trading_date, profile, time_range, FLAGS.min_order_interval_ms))
  sys.stdout.flush()
  trading_dates = abase.get_trading_dates(trading_date)
  launch_multi(trading_dates=trading_dates,
               time_range=time_range,
               profile=profile,
               thresholds_bps=FLAGS.threshold_bps.split(","),
               leans_bps=FLAGS.lean_bps.split(","),
               stacks=list(numpy.array(FLAGS.stacks.split(",")).astype(int)))


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()

  flags.DEFINE_string('stacks', '2,8,16', '')

  flags.DEFINE_string('profile', None, '')

  flags.DEFINE_string('threshold_bps', '0', '')

  flags.DEFINE_string('lean_bps', '0', '')

  flags.DEFINE_float('min_order_interval_ms', 200, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
